import cv2
import numpy as np
import os
import random

def poisson_blending(foreground, background, mask, center):
    """
    使用泊松融合将前景图像与背景图像结合。
    :param foreground: 前景图像
    :param background: 背景图像
    :param mask: 掩膜图像（单通道）
    :param center: 融合区域的中心点
    :return: 融合后的图像
    """
    blended = cv2.seamlessClone(foreground, background, mask, center, cv2.NORMAL_CLONE)
    return blended

def process_images(folder_a, folder_b, mask_folder, output_folder):
    """
    遍历前景图像，并使用对应的掩膜进行泊松融合。
    :param folder_a: 前景图像文件夹路径
    :param folder_b: 背景图像文件夹路径
    :param mask_folder: 掩膜文件夹路径
    :param output_folder: 输出图像文件夹路径
    """
    # 确保输出文件夹存在
    os.makedirs(output_folder, exist_ok=True)

    # 获取文件夹中的所有文件名
    files_a = sorted(os.listdir(folder_a))
    files_b = sorted(os.listdir(folder_b))

    if not files_a or not files_b:
        print("前景或背景文件夹为空，无法进行融合。")
        return

    # 遍历前景文件，并使用对应的掩膜文件
    for i, file_a in enumerate(files_a):
        file_b = random.choice(files_b)  # 随机选择一个背景图像
        mask_path = os.path.join(mask_folder, file_a)  # 掩膜文件路径

        # 加载前景、背景和掩膜图像
        foreground = cv2.imread(os.path.join(folder_a, file_a))
        background = cv2.imread(os.path.join(folder_b, file_b))
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

        if foreground is None or background is None:
            print(f"无法加载图像：{file_a} 或 {file_b}")
            continue

        if mask is None:
            print(f"掩膜文件缺失，跳过：{file_a}")
            continue

        # 确保前景、背景和掩膜大小一致
        background = cv2.resize(background, (foreground.shape[1], foreground.shape[0]))
        mask = cv2.resize(mask, (foreground.shape[1], foreground.shape[0]))

        # 确保掩膜是二值图像
        _, mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)

        # 定义融合区域的中心点（背景图像的中心）
        center = (background.shape[1] // 2, background.shape[0] // 2)

        # 使用泊松融合
        result = poisson_blending(foreground, background, mask, center)

        # 保存结果图像
        output_path = os.path.join(output_folder, f"Img_{i}.png")
        cv2.imwrite(output_path, result)

        print(f"已保存融合图像：{output_path}")

# 指定文件夹路径
folder_a = "rawA"  # 前景图像文件夹路径
folder_b = "rawB"  # 背景图像文件夹路径
mask_folder = "msk"  # 掩膜文件夹路径
output_folder = "C1"  # 输出图像文件夹路径

# 执行融合
process_images(folder_a, folder_b, mask_folder, output_folder)
